// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/exec/spill.h>
#include <pollux/plan/plan_builder.h>
#include <pollux/common/file/temp_directory_path.h>
#include <pollux/functions/lib/aggregates/tests/utils/AggregationTestBase.h>
#include <pollux/functions/lib/window/tests/WindowTestBase.h>

using namespace kumo::pollux::exec::test;
using namespace kumo::pollux::functions::aggregate::test;
using namespace kumo::pollux::window::test;

namespace kumo::pollux::aggregate::test {

namespace {

class ArbitraryTest : public AggregationTestBase {};

TEST_F(ArbitraryTest, noNulls) {
  // Create vectors without nulls because DuckDB's "first" aggregate does not
  // ignore them.
  const int32_t size = 10'000;
  auto vectors = {make_row_vector(
      {make_flat_vector<int64_t>(size, [](vector_size_t row) { return row; }),
       make_flat_vector<int8_t>(size, [](vector_size_t row) { return row; }),
       make_flat_vector<int16_t>(size, [](vector_size_t row) { return row; }),
       make_flat_vector<int32_t>(size, [](vector_size_t row) { return row; }),
       make_flat_vector<int64_t>(size, [](vector_size_t row) { return row; }),
       make_flat_vector<float>(size, [](vector_size_t row) { return row; }),
       make_flat_vector<double>(size, [](vector_size_t row) { return row; })})};
  createDuckDbTable(vectors);

  std::vector<std::string> aggregates = {
      "arbitrary(c1)",
      "arbitrary(c2)",
      "any_value(c3)",
      "arbitrary(c4)",
      "arbitrary(c5)",
      "any_value(c6)"};

  // We do not test with TableScan because having two input splits makes the
  // result non-deterministic.
  // Global aggregation.
  testAggregations(
      vectors,
      {},
      aggregates,
      "SELECT first(c1), first(c2), first(c3), first(c4), first(c5), first(c6) FROM tmp");

  // Group by aggregation.
  testAggregations(
      [&](PlanBuilder& builder) {
        builder.values(vectors).project(
            {"c0 % 10", "c1", "c2", "c3", "c4", "c5", "c6"});
      },
      {"p0"},
      aggregates,
      "SELECT c0 % 10, first(c1), first(c2), first(c3), first(c4), first(c5), first(c6) FROM tmp GROUP BY 1");

  // encodings: use filter to wrap aggregation inputs in a dictionary.
  testAggregations(
      [&](PlanBuilder& builder) {
        builder.values(vectors)
            .filter("c0 % 2 = 0")
            .project({"c0 % 10", "c1", "c2", "c3", "c4", "c5", "c6"});
      },
      {"p0"},
      aggregates,
      "SELECT c0 % 10, first(c1), first(c2), first(c3), first(c4), first(c5), first(c6) FROM tmp WHERE c0 % 2 = 0 GROUP BY 1");

  testAggregations(
      [&](PlanBuilder& builder) {
        builder.values(vectors).filter("c0 % 2 = 0");
      },
      {},
      aggregates,
      "SELECT first(c1), first(c2), first(c3), first(c4), first(c5), first(c6) FROM tmp WHERE c0 % 2 = 0");
}

TEST_F(ArbitraryTest, nulls) {
  auto vectors = {
      make_row_vector(
          {make_nullable_flat_vector<int32_t>({1, 1, 2, 2, 3, 3}),
           make_nullable_flat_vector<int64_t>(
               {std::nullopt, std::nullopt, std::nullopt, 4, std::nullopt, 5}),
           make_nullable_flat_vector<double>({
               std::nullopt,
               0.50,
               std::nullopt,
               std::nullopt,
               0.25,
               std::nullopt,
           }),
           make_null_constant(TypeKind::UNKNOWN, 6)}),
  };

  // We do not test with TableScan because having two input splits makes the
  // result non-deterministic. Also, unknown type is not supported in Writer
  // yet. Global aggregation.
  testAggregations(
      vectors,
      {},
      {"arbitrary(c1)", "arbitrary(c2)", "arbitrary(c3)"},
      "SELECT * FROM( VALUES (4, 0.50, NULL)) AS t");

  // Group by aggregation.
  testAggregations(
      vectors,
      {"c0"},
      {"arbitrary(c1)", "arbitrary(c2)", "arbitrary(c3)"},
      "SELECT * FROM(VALUES (1, NULL, 0.50, NULL), (2, 4, NULL, NULL), (3, 5, 0.25, NULL)) AS t");
}

TEST_F(ArbitraryTest, varchar) {
  auto rowType = ROW({"c0", "c1"}, {INTEGER(), VARCHAR()});
  auto vectors = makeVectors(rowType, 1000, 10);
  createDuckDbTable(vectors);

  // We do not test with TableScan because having two input splits makes the
  // result non-deterministic.
  testAggregations(
      [&](PlanBuilder& builder) {
        builder.values(vectors).project({"c0 % 11", "c1"});
      },
      {"p0"},
      {"arbitrary(c1)"},
      "SELECT c0 % 11, first(c1) FROM tmp WHERE c1 IS NOT NULL GROUP BY 1");

  testAggregations(
      vectors,
      {},
      {"arbitrary(c1)"},
      "SELECT first(c1) FROM tmp WHERE c1 IS NOT NULL");

  // encodings: use filter to wrap aggregation inputs in a dictionary.
  testAggregations(
      [&](PlanBuilder& builder) {
        builder.values(vectors).filter("c0 % 2 = 0").project({"c0 % 11", "c1"});
      },
      {"p0"},
      {"arbitrary(c1)"},
      "SELECT c0 % 11, first(c1) FROM tmp WHERE c0 % 2 = 0 AND c1 IS NOT NULL GROUP BY 1");

  testAggregations(
      [&](PlanBuilder& builder) {
        builder.values(vectors).filter("c0 % 2 = 0");
      },
      {},
      {"arbitrary(c1)"},
      "SELECT first(c1) FROM tmp WHERE c0 % 2 = 0 AND c1 IS NOT NULL");
}

TEST_F(ArbitraryTest, varcharConstAndNulls) {
  auto vectors = {make_row_vector({
      make_flat_vector<int32_t>(100, [](auto row) { return row % 7; }),
      make_constant("apple", 100),
      make_null_constant(TypeKind::VARCHAR, 100),
  })};

  createDuckDbTable(vectors);

  testAggregations(
      vectors,
      {},
      {"arbitrary(c1)", "arbitrary(c2)"},
      "SELECT first(c1), first(c2) FROM tmp");

  testAggregations(
      vectors,
      {"c0"},
      {"arbitrary(c1)", "arbitrary(c2)"},
      "SELECT c0, first(c1), first(c2) FROM tmp group by c0");
}

TEST_F(ArbitraryTest, numericConstAndNulls) {
  auto vectors = {make_row_vector(
      {make_flat_vector<int32_t>(100, [](auto row) { return row % 7; }),
       make_constant(11, 100),
       make_null_constant(TypeKind::BIGINT, 100)})};

  createDuckDbTable(vectors);

  testAggregations(
      vectors,
      {},
      {"arbitrary(c1)", "arbitrary(c2)"},
      "SELECT first(c1), first(c2) FROM tmp");

  testAggregations(
      vectors,
      {"c0"},
      {"arbitrary(c1)", "arbitrary(c2)"},
      "SELECT c0, first(c1), first(c2) FROM tmp group by c0");
}

TEST_F(ArbitraryTest, boolean) {
  auto data = make_row_vector({
      // Grouping key.
      make_flat_vector<int64_t>({1, 1, 2, 2, 3, 3, 4, 4}),
      // Input values: 'constant' within groups.
      make_nullable_flat_vector<bool>(
          {true,
           true,
           false,
           false,
           std::nullopt,
           std::nullopt,
           std::nullopt,
           false}),
      make_constant<bool>(std::nullopt, 8),
  });

  auto expectedResult = make_row_vector({
      make_flat_vector<int64_t>({1, 2, 3, 4}),
      make_nullable_flat_vector<bool>({true, false, std::nullopt, false}),
  });

  testAggregations({data}, {"c0"}, {"arbitrary(c1)"}, {expectedResult});

  // Global aggregation.
  auto plan = PlanBuilder()
                  .values({data})
                  .singleAggregation({}, {"arbitrary(c1)"})
                  .planNode();

  assertQuery(plan, "SELECT true");

  testAggregations({data}, {}, {"arbitrary(c2)"}, "SELECT null");
}

TEST_F(ArbitraryTest, timestamp) {
  auto data = make_row_vector({
      // Grouping key.
      make_flat_vector<int64_t>({1, 1, 2, 2, 3, 3, 4, 4}),
      // Input values: constant within groups: 100.1, 100.1, 200.2, 200.2, etc.
      make_nullable_flat_vector<Timestamp>(
          {Timestamp(100, 1),
           Timestamp(100, 1),
           Timestamp(200, 2),
           Timestamp(200, 2),
           std::nullopt,
           std::nullopt,
           std::nullopt,
           Timestamp(100, 4)}),
      make_constant<Timestamp>(std::nullopt, 8),
  });

  auto expectedResult = make_row_vector({
      make_flat_vector<int64_t>({1, 2, 3, 4}),
      make_nullable_flat_vector<Timestamp>(
          {Timestamp(100, 1),
           Timestamp(200, 2),
           std::nullopt,
           Timestamp(100, 4)}),
  });

  testAggregations({data}, {"c0"}, {"arbitrary(c1)"}, {expectedResult});

  // Global aggregation.
  auto plan = PlanBuilder()
                  .values({data})
                  .singleAggregation({}, {"arbitrary(c1)"})
                  .planNode();

  auto result = readSingleValue(plan);
  ASSERT_TRUE(!result.isNull());
  ASSERT_EQ(result.kind(), TypeKind::TIMESTAMP);

  auto timestamp = result.value<Timestamp>();
  ASSERT_EQ(timestamp, Timestamp(100, 1));

  testAggregations({data}, {}, {"arbitrary(c2)"}, "SELECT null");
}

TEST_F(ArbitraryTest, date) {
  auto data = make_row_vector({
      // Grouping key.
      make_flat_vector<int64_t>({1, 1, 2, 2, 3, 3, 4, 4}),
      // Input values: constant within groups.
      make_nullable_flat_vector<int32_t>(
          {125, 125, 126, 126, std::nullopt, std::nullopt, std::nullopt, 128},
          DATE()),
      make_constant<Timestamp>(std::nullopt, 8),
  });

  auto expectedResult = make_row_vector({
      make_flat_vector<int64_t>({1, 2, 3, 4}),
      make_nullable_flat_vector<int32_t>({125, 126, std::nullopt, 128}, DATE()),
  });

  testAggregations({data}, {"c0"}, {"arbitrary(c1)"}, {expectedResult});

  // Global aggregation.
  auto plan = PlanBuilder()
                  .values({data})
                  .singleAggregation({}, {"arbitrary(c1)"})
                  .planNode();

  auto result = readSingleValue(plan);
  ASSERT_TRUE(!result.isNull());
  ASSERT_EQ(result.kind(), TypeKind::INTEGER);

  auto date = result.value<int32_t>();
  ASSERT_EQ(date, 125);

  testAggregations({data}, {}, {"arbitrary(c2)"}, "SELECT null");
}

TEST_F(ArbitraryTest, interval) {
  auto data = make_row_vector({
      // Grouping key.
      make_flat_vector<int64_t>({1, 1, 2, 2, 3, 3, 4, 4}),
      // Input values: constant within groups.
      make_nullable_flat_vector<int64_t>(
          {125, 125, 126, 126, std::nullopt, std::nullopt, std::nullopt, 128},
          INTERVAL_DAY_TIME()),
      make_constant<Timestamp>(std::nullopt, 8),
  });

  auto expectedResult = make_row_vector({
      make_flat_vector<int64_t>({1, 2, 3, 4}),
      make_nullable_flat_vector<int64_t>(
          {125, 126, std::nullopt, 128}, INTERVAL_DAY_TIME()),
  });

  testAggregations({data}, {"c0"}, {"arbitrary(c1)"}, {expectedResult});

  // Global aggregation.
  auto plan = PlanBuilder()
                  .values({data})
                  .singleAggregation({}, {"arbitrary(c1)"})
                  .planNode();

  auto interval = readSingleValue(plan);
  ASSERT_EQ(interval.value<int64_t>(), 125);

  testAggregations({data}, {}, {"arbitrary(c2)"}, "SELECT null");
}

TEST_F(ArbitraryTest, longDecimal) {
  auto data = make_row_vector({// Grouping key.
                             make_flat_vector<int64_t>({1, 1, 2, 2, 3, 3, 4, 4}),
                             make_nullable_flat_vector<int128_t>(
                                 {HugeInt::build(10, 100),
                                  HugeInt::build(10, 100),
                                  HugeInt::build(10, 200),
                                  HugeInt::build(10, 200),
                                  std::nullopt,
                                  std::nullopt,
                                  std::nullopt,
                                  HugeInt::build(10, 400)},
                                 DECIMAL(38, 8))});

  auto expectedResult = make_row_vector({
      make_flat_vector<int64_t>({1, 2, 3, 4}),
      make_nullable_flat_vector<int128_t>(
          {HugeInt::build(10, 100),
           HugeInt::build(10, 200),
           std::nullopt,
           HugeInt::build(10, 400)},
          DECIMAL(38, 8)),
  });

  testAggregations({data}, {"c0"}, {"arbitrary(c1)"}, {expectedResult});
}

TEST_F(ArbitraryTest, shortDecimal) {
  auto data = make_row_vector({// Grouping key.
                             make_flat_vector<int64_t>({1, 1, 2, 2, 3, 3, 4, 4}),
                             make_nullable_flat_vector<int64_t>(
                                 {10000000000000000,
                                  10000000000000000,
                                  20000000000000000,
                                  20000000000000000,
                                  std::nullopt,
                                  std::nullopt,
                                  std::nullopt,
                                  40000000000000000},
                                 DECIMAL(15, 2))});

  auto expectedResult = make_row_vector({
      make_flat_vector<int64_t>({1, 2, 3, 4}),
      make_nullable_flat_vector<int64_t>(
          {10000000000000000,
           20000000000000000,
           std::nullopt,
           40000000000000000},
          DECIMAL(15, 2)),
  });

  testAggregations({data}, {"c0"}, {"arbitrary(c1)"}, {expectedResult});
}

class ArbitraryWindowTest : public WindowTestBase {};

TEST_F(ArbitraryWindowTest, basic) {
  auto data = make_row_vector(
      {make_flat_vector<int64_t>({1, 2, 3, 4, 5}),
       make_array_vector<double>({{1.0}, {2.0}, {3.0}, {4.0}, {5.0}}),
       make_flat_vector<bool>({false, false, false, false, false})});

  auto expected = make_row_vector(
      {make_flat_vector<int64_t>({1, 2, 3, 4, 5}),
       make_array_vector<double>({{1.0}, {2.0}, {3.0}, {4.0}, {5.0}}),
       make_flat_vector<bool>({false, false, false, false, false}),
       make_flat_vector<int64_t>({1, 1, 1, 1, 1})});
  window::test::WindowTestBase::testWindowFunction(
      {data},
      "arbitrary(c0)",
      "partition by c2 order by c0",
      "range between unbounded preceding and current row",
      expected);

  expected = make_row_vector(
      {make_flat_vector<int64_t>({1, 2, 3, 4, 5}),
       make_array_vector<double>({{1.0}, {2.0}, {3.0}, {4.0}, {5.0}}),
       make_flat_vector<bool>({false, false, false, false, false}),
       make_array_vector<double>({{1.0}, {1.0}, {1.0}, {1.0}, {1.0}})});
  window::test::WindowTestBase::testWindowFunction(
      {data},
      "arbitrary(c1)",
      "partition by c2 order by c0",
      "range between unbounded preceding and current row",
      expected);
}

TEST_F(ArbitraryTest, spilling) {
  auto data = make_row_vector(
      {make_flat_vector<float>({0.1, 0.2, 0.3, 0.4, 0.5, 0.6}),
       make_nullable_flat_vector<int64_t>({1, 2, 3, 4, 5, 6})});
  auto expected = make_row_vector(
      {make_nullable_flat_vector<int64_t>({1, 2, 3, 4, 5, 6}),
       make_nullable_flat_vector<float>({0.1, 0.2, 0.3, 0.4, 0.5, 0.6})});

  auto plan = PlanBuilder()
                  .values({data})
                  .singleAggregation({"c1"}, {"arbitrary(c0)"})
                  .planNode();

  std::shared_ptr<TempDirectoryPath> spillDirectory;
  AssertQueryBuilder builder(plan);

  exec::TestScopedSpillInjection scopedSpillInjection(100);
  spillDirectory = exec::test::TempDirectoryPath::create();
  builder.spillDirectory(spillDirectory->getPath())
      .config(core::QueryConfig::kSpillEnabled, "true")
      .config(core::QueryConfig::kAggregationSpillEnabled, "true")
      .config(core::QueryConfig::kSpillNumPartitionBits, "0");

  auto result = builder.maxDrivers(2).copyResults(pool_.get());
  ::kumo::pollux::test::assertEqualVectors(expected, result);
}

} // namespace
} // namespace kumo::pollux::aggregate::test
